package com.fuyun.rpc.client.rateLimiting;

import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 令牌桶限流器
 */
public class TokenBucketRateLimiter {

	// 令牌桶容量
	private final int capacity;

	// 每秒产生的令牌数
	private final int refillTokensPerSecond;

	// 当前令牌数
	private final AtomicInteger tokens = new AtomicInteger(0);

	// 定时任务执行器
	private final ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(1);

	public TokenBucketRateLimiter(int capacity, int refillTokensPerSecond) {
		this.capacity = capacity;
		this.refillTokensPerSecond = refillTokensPerSecond;
		this.tokens.set(capacity);

		// 定时补充令牌
		scheduler.scheduleAtFixedRate(this::refill, 1, 1, TimeUnit.SECONDS);
	}

	/**
	 * 补充令牌
	 */
	private void refill() {
		int currentTokens = tokens.get();
		int newTokens = Math.min(capacity, currentTokens + refillTokensPerSecond);
		tokens.compareAndSet(currentTokens, newTokens);
	}

	/**
	 * 尝试获取令牌
	 * @return 是否获取成功
	 */
	public boolean tryAcquire() {
		return tryAcquire(1);
	}

	/**
	 * 尝试获取指定数量的令牌
	 * @param numTokens 需要的令牌数
	 * @return 是否获取成功
	 */
	public boolean tryAcquire(int numTokens) {
		while (true) {
			int currentTokens = tokens.get();
			if (currentTokens < numTokens) {
				return false;
			}

			if (tokens.compareAndSet(currentTokens, currentTokens - numTokens)) {
				return true;
			}
		}
	}

	/**
	 * 获取当前剩余令牌数
	 * @return 剩余令牌数
	 */
	public int getAvailableTokens() {
		return tokens.get();
	}

	/**
	 * 关闭限流器
	 */
	public void shutdown() {
		scheduler.shutdown();
	}
}